"""
中文文本智能分块工具
专门针对中文段落、层级结构和列表内容优化
"""
from typing import List, Dict, Tuple
import re


class ChineseTextChunker:
    """中文文本智能分块器"""
    
    def __init__(self, 
                 max_chunk_size: int = 1200,
                 min_chunk_size: int = 200,
                 overlap_size: int = 100):
        """
        Args:
            max_chunk_size: 最大块大小（字符数）
            min_chunk_size: 最小块大小（字符数）
            overlap_size: 块之间的重叠大小
        """
        self.max_chunk_size = max_chunk_size
        self.min_chunk_size = min_chunk_size
        self.overlap_size = overlap_size
        
        # 标题模式定义（按优先级排序）
        self.heading_patterns = [
            # 注意：多级小数点编号需要特殊处理，不能用简单的正则匹配层级
            # 所以我们用更灵活的方法：先识别是否是标题，再判断层级
        ]
    
    def get_line_heading_level(self, line: str) -> Tuple[int, str]:
        """
        判断行的标题层级（改进版：支持多级小数点编号）
        
        层级判断规则：
        - 0级：普通文本
        - 1级：单个数字编号（1. 2. 一、二、第一章等）
        - 2级：两段编号（1.1 8.1 (1) 等）
        - 3级：三段编号（1.1.1 8.1.1 等）或列表项（a) b) ① 等）
        - 4级：四段及以上编号（1.1.1.1 等）或更深层列表
        
        Args:
            line: 文本行
            
        Returns:
            Tuple[int, str]: (层级, 标题文本)，如果不是标题返回(0, "")
        """
        line = line.strip()
        if not line or len(line) > 100:  # 标题一般不会太长
            return 0, ""
        
        # 1. 检查多级小数点编号（如 8.1, 8.1.1, 8.1.2.1）
        decimal_pattern = r'^(\d+(?:\.\d+)*)[.、\s]+'
        decimal_match = re.match(decimal_pattern, line)
        if decimal_match:
            number_part = decimal_match.group(1)
            dot_count = number_part.count('.')
            
            # 根据小数点数量判断层级
            # 8 -> 1级, 8.1 -> 2级, 8.1.1 -> 3级, 8.1.1.1 -> 4级
            level = min(dot_count + 1, 4)  # 最多4级
            return level, line
        
        # 2. 检查中文编号（一、二、三、第一章、第二节）
        chinese_num_patterns = [
            (r'^[一二三四五六七八九十百千]+[、.]\s*.{1,50}$', 1),
            (r'^第[一二三四五六七八九十百千]+[章节条款]\s*.{1,50}$', 1),
        ]
        for pattern, level in chinese_num_patterns:
            if re.match(pattern, line):
                return level, line
        
        # 3. 检查括号编号（(1) (2)）- 视为2级标题
        if re.match(r'^\([1-9]\d*\)\s*.{1,50}$', line):
            return 2, line
        
        # 4. 检查列表项（a) b) c) 或 ① ② ③ 或 1) 2) 3)）- 视为3级标题
        list_patterns = [
            r'^[a-z]\)\s*.{1,200}$',                          # a) b) c)
            r'^[A-Z]\)\s*.{1,200}$',                          # A) B) C)
            r'^[①②③④⑤⑥⑦⑧⑨⑩⑪⑫⑬⑭⑮⑯⑰⑱⑲⑳]\s*.{1,200}$',  # ① ② ③
            r'^[1-9]\d*\)\s*.{1,200}$',                       # 1) 2) 3)
        ]
        for pattern in list_patterns:
            if re.match(pattern, line):
                return 3, line
        
        # 5. 检查罗马数字（I. II. III.）- 视为1级标题
        if re.match(r'^[IVX]+[.、]\s*.{1,50}$', line):
            return 1, line
        
        # 不是任何已知的标题格式
        return 0, ""
    
    def is_sentence_end(self, text: str) -> bool:
        """
        判断文本是否以句子结束符结尾
        
        Args:
            text: 文本
            
        Returns:
            bool: 是否为句子结束
        """
        sentence_endings = ['。', '！', '？', '；', '.', '!', '?', ';']
        return any(text.rstrip().endswith(end) for end in sentence_endings)
    
    def find_best_split_point(self, lines: List[str], start: int, max_end: int) -> int:
        """
        在指定范围内找到最佳分割点（改进版：支持多级标题）
        优先考虑：1. 一级标题 2. 二级标题 3. 三级标题 4. 句子结束 5. 段落结束
        
        Args:
            lines: 所有文本行
            start: 起始行
            max_end: 最大结束行
            
        Returns:
            int: 最佳分割点行号
        """
        best_split = max_end
        best_priority = 999  # 数字越小优先级越高
        
        # 从后往前找，记录最佳分割点
        for i in range(max_end - 1, start, -1):
            level, _ = self.get_line_heading_level(lines[i])
            
            # 优先级：1级标题 > 2级标题 > 3级标题 > 句子结束
            current_priority = 999
            
            if level == 1:
                current_priority = 1  # 最高优先级
            elif level == 2:
                current_priority = 2
            elif level == 3:
                current_priority = 3
            elif i > 0 and self.is_sentence_end(lines[i - 1]):
                current_priority = 4
            
            # 更新最佳分割点（取优先级最高且最靠后的位置）
            if current_priority < best_priority:
                best_split = i
                best_priority = current_priority
                
                # 如果找到1级或2级标题，就可以直接返回了
                if level in [1, 2]:
                    return best_split
        
        return best_split
    
    def should_keep_together(self, lines: List[str], start: int, end: int) -> bool:
        """
        判断一段文本是否应该保持在一起（不分割）
        主要用于识别紧密关联的列表项（改进版：支持多级结构）
        
        规则：
        - 同一小节下的列表项（3级标题）应该保持在一起
        - 除非内容太长超过阈值
        
        Args:
            lines: 文本行列表
            start: 起始行
            end: 结束行
            
        Returns:
            bool: 是否应该保持在一起
        """
        if end <= start:
            return False
        
        # 统计各层级标题的数量
        level_counts = {1: 0, 2: 0, 3: 0, 4: 0}
        for i in range(start, min(end, len(lines))):
            level, _ = self.get_line_heading_level(lines[i])
            if level > 0:
                level_counts[level] += 1
        
        # 如果只有3级或4级标题（列表项），尝试保持在一起
        if level_counts[3] >= 2 and level_counts[1] == 0 and level_counts[2] == 0:
            total_chars = sum(len(lines[i]) for i in range(start, min(end, len(lines))))
            # 只有在总字符数不太大的情况下才保持在一起
            if total_chars < self.max_chunk_size * 0.9:
                return True
        
        return False
    
    def create_chunk(self, lines: List[str], start: int, end: int, 
                    chunk_id: int, parent_title: str = "") -> Dict:
        """
        创建一个文本块（改进版：支持多级标题提取）
        
        Args:
            lines: 所有文本行
            start: 起始行
            end: 结束行
            chunk_id: 块ID
            parent_title: 父标题
            
        Returns:
            Dict: 块信息，包含 chunk_id, title, text, char_count, line_range, level
        """
        chunk_lines = lines[start:end]
        chunk_text = '\n'.join(chunk_lines)
        
        # 提取块的主标题（优先选择层级最高的标题）
        title = parent_title
        title_level = 0
        
        for line in chunk_lines:
            level, heading_text = self.get_line_heading_level(line)
            # 选择层级最高（数字最小）的标题
            if level > 0 and (title_level == 0 or level < title_level):
                title = heading_text
                title_level = level
                # 如果找到1级或2级标题，就不用继续找了
                if level <= 2:
                    break
        
        # 如果没有找到标题，使用第一个非空行
        if not title and chunk_lines:
            for line in chunk_lines:
                if line.strip():
                    title = line.strip()[:30] + "..." if len(line.strip()) > 30 else line.strip()
                    break
        
        return {
            "chunk_id": chunk_id,
            "title": title if title else f"块{chunk_id}",
            "text": chunk_text,
            "char_count": len(chunk_text),
            "line_range": (start, end),
            "level": title_level  # 添加层级信息
        }
    
    def split_text_into_chunks(self, text: str) -> List[Dict]:
        """
        将文本智能分割为块（改进版：支持多级标题结构）
        
        策略：
        1. 在1级或2级标题处优先分割
        2. 保持3级标题（列表项）与其所属章节在一起
        3. 当块太大时，在句子边界分割
        4. 避免过小的块（合并或过滤）
        
        Args:
            text: 要分割的文本
            
        Returns:
            List[Dict]: 分块结果列表
        """
        lines = text.split('\n')
        chunks = []
        
        current_start = 0
        current_size = 0
        parent_title = ""
        last_major_heading_level = 0  # 记录最近的主标题层级
        
        i = 0
        while i < len(lines):
            line = lines[i]
            level, heading = self.get_line_heading_level(line)
            
            # 遇到1级或2级标题时考虑分割
            if level in [1, 2]:
                # 如果已经累积了足够的内容，先创建一个块
                if current_size > self.min_chunk_size and current_start < i:
                    chunk = self.create_chunk(lines, current_start, i, len(chunks), parent_title)
                    if chunk["char_count"] > 50:  # 过滤太小的块
                        chunks.append(chunk)
                    current_start = i
                    current_size = 0
                
                # 更新父标题和层级
                parent_title = heading
                last_major_heading_level = level
            
            current_size += len(line)
            
            # 检查是否需要分割（块太大）
            if current_size >= self.max_chunk_size:
                # 查找最佳分割点
                best_split = self.find_best_split_point(lines, current_start, i + 1)
                
                # 检查是否应该保持在一起（列表项聚合）
                if not self.should_keep_together(lines, current_start, best_split):
                    chunk = self.create_chunk(lines, current_start, best_split, len(chunks), parent_title)
                    if chunk["char_count"] > 50:
                        chunks.append(chunk)
                    
                    # 添加少量重叠以保持上下文连贯
                    overlap_start = max(current_start, best_split - 2)  # 重叠2行
                    current_start = overlap_start
                    current_size = sum(len(lines[j]) for j in range(overlap_start, i + 1))
                else:
                    # 如果必须保持在一起但又太大，直接分割（避免无限增长）
                    if current_size > self.max_chunk_size * 1.5:
                        chunk = self.create_chunk(lines, current_start, i + 1, len(chunks), parent_title)
                        if chunk["char_count"] > 50:
                            chunks.append(chunk)
                        current_start = i + 1
                        current_size = 0
            
            i += 1
        
        # 处理剩余内容
        if current_start < len(lines):
            chunk = self.create_chunk(lines, current_start, len(lines), len(chunks), parent_title)
            if chunk["char_count"] > 50:
                chunks.append(chunk)
        
        # 重新分配chunk_id
        for idx, chunk in enumerate(chunks):
            chunk["chunk_id"] = idx
        
        return chunks


def split_text_into_chunks(text: str, max_chunk_size: int = 1200) -> List[Dict]:
    """
    将文本分割为块（对外接口，保持兼容性）
    
    Args:
        text: 要分割的文本
        max_chunk_size: 最大块大小（字符数）
        
    Returns:
        List[Dict]: 分块结果列表
    """
    chunker = ChineseTextChunker(max_chunk_size=max_chunk_size)
    return chunker.split_text_into_chunks(text)


if __name__ == "__main__":
    # 测试代码
    test_text = """2.业务数据安全保障

(1) 数据收集

a) 查验企业是否根据业务需要明确收集数据的目的和用途，仅收集与业务需求相关的最小数据集，避免过度收集不必要的数据，并向客户明确告知数据收集目的、途径和范围，获得客户授权后再进行数据收集。

b) 查验企业是否就涉个人数据收集保障个人的知情权和选择权，遵守GB/T 35273-2020第5章的要求，包括收集个人数据需合法、保持最小必要性、可自主选择，并获得授权同意。

(2) 数据存储

a) 查验企业是否采取必要的技术措施保障数据存储安全。

b) 查验企业是否建立完善的数据备份机制。"""
    
    chunker = ChineseTextChunker(max_chunk_size=400)
    chunks = chunker.split_text_into_chunks(test_text)
    
    print(f"总共分割为 {len(chunks)} 个块：\n")
    for chunk in chunks:
        print(f"块ID: {chunk['chunk_id']}")
        print(f"标题: {chunk['title']}")
        print(f"字符数: {chunk['char_count']}")
        print(f"内容预览: {chunk['text'][:100]}...")
        print("-" * 80)

